import torch
import torch.nn as nn
import numpy as np
from ..module import build_backbone, DETECTOR
from ..loss import loss_dict


@DETECTOR.register_module
class BaseDetector(nn.Module):

    def __init__(self, backbone_cfg, stride, feat_size, fc_channels, losses, mode='train'):
        super(BaseDetector, self).__init__()
        self.backbone = build_backbone(backbone_cfg)
        self.stride = stride
        self.gap = nn.AvgPool2d(feat_size[1])
        self.fcs = nn.ModuleList()
        for i in range(len(fc_channels)):
            if i != 0:
                self.fcs.append(nn.Linear(fc_channels[i-1], fc_channels[i]))
            else:
                self.fcs.append(nn.Linear(feat_size[0], fc_channels[i]))
        self.losses = []
        for loss in losses:
            self.losses.append(loss_dict[loss])
        self.mode = mode
        self.loss_rds = np.zeros([200])
        self.loss_rd_ind = 0
        self.loss_rd_sum = 0.

    def forward(self, feature, label=None):
        backbone_feature = self.backbone(feature)
        x = self.gap(backbone_feature)
        x = x.view([len(x), -1])
        for fc in self.fcs:
            x = fc(x)
        cls = x
        if self.mode == 'train':
            loss_func = self.losses[0]
            cls_loss = loss_func(cls, label)
            loss_rd = cls_loss.cpu().data.numpy()
            self.loss_rds[self.loss_rd_ind] = loss_rd
            self.loss_rd_ind = (self.loss_rd_ind+1) % 200
            self.loss_rd_sum += loss_rd
            self.loss_rd_sum -= self.loss_rds[self.loss_rd_ind]
            return cls_loss, {'cls_loss': self.loss_rd_sum/200}
        else:
            return cls, self.bbox_predict(backbone_feature, cls)

    def init_weights(self):
        self.backbone.init_weights()
        for fc in self.fcs:
            torch.nn.init.normal_(fc.weight, std=0.01)
            torch.nn.init.constant_(fc.bias, 0)

    def bbox_predict(self, backbone_feature, cls, threshold=0.05):
        clses = cls.max(1)[1]
        h_size = backbone_feature.size()
        weights = []
        for i in range(h_size[0]):
            weights.append(self.fcs[-1].weight[clses[i], :])
        weights = torch.stack(weights, dim=0)
        weights = weights.unsqueeze(2).unsqueeze(3).repeat(1, 1, h_size[2], h_size[3])

        heat_map = (backbone_feature * weights).sum(1)

        max_v = heat_map.max()
        min_v = heat_map.min()
        horizen_line = heat_map.max(1)[0]
        horizen_line = (horizen_line - min_v) / (max_v - min_v + 1e-8)
        vertical_line = heat_map.max(2)[0]
        vertical_line = (vertical_line - min_v) / (max_v - min_v + 1e-8)

        bboxes = torch.zeros(size=(heat_map.shape[0], 4), dtype=torch.int32, device=heat_map.device)
        for c in range(h_size[0]):

            x1 = 0
            x2 = 0
            y1 = 0
            y2 = 0
            for i in range(h_size[3]):
                if horizen_line[c, i] > threshold:
                    if x1 == 0:
                        x1 = i
                    x2 = i
            for i in range(h_size[2]):
                if vertical_line[c, i] > threshold:
                    if y1 == 0:
                        y1 = i
                    y2 = i
            bboxes[c, 0] = x1
            bboxes[c, 1] = y1
            bboxes[c, 2] = x2
            bboxes[c, 3] = y2
        bboxes = bboxes * self.stride
        return bboxes